import os
import pandas as pd
import pickle
from datetime import datetime, timedelta
import json

class DataCache:
    """数据缓存管理器，用于缓存股票数据，减少API调用"""
    
    def __init__(self, cache_dir="./cache"):
        """初始化缓存管理器"""
        self.cache_dir = cache_dir
        # 确保缓存目录存在
        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
        
        # 缓存配置
        self.config_file = os.path.join(cache_dir, "cache_config.json")
        self.load_config()
    
    def load_config(self):
        """加载缓存配置"""
        if os.path.exists(self.config_file):
            with open(self.config_file, 'r', encoding='utf-8') as f:
                self.config = json.load(f)
        else:
            # 默认配置
            self.config = {
                "stock_list_ttl": 86400,  # 股票列表缓存1天
                "daily_data_ttl": 3600,   # 日线数据缓存1小时
                "financial_data_ttl": 86400,  # 财务数据缓存1天
                "last_update": {}
            }
            self.save_config()
    
    def save_config(self):
        """保存缓存配置"""
        with open(self.config_file, 'w', encoding='utf-8') as f:
            json.dump(self.config, f, ensure_ascii=False, indent=2)
    
    def get_cache_path(self, cache_type, key=None):
        """获取缓存文件路径"""
        if key:
            return os.path.join(self.cache_dir, f"{cache_type}_{key}.pkl")
        return os.path.join(self.cache_dir, f"{cache_type}.pkl")
    
    def is_cache_valid(self, cache_type, key=None):
        """检查缓存是否有效"""
        cache_path = self.get_cache_path(cache_type, key)
        
        # 如果缓存文件不存在，则无效
        if not os.path.exists(cache_path):
            return False
        
        # 检查缓存是否过期
        cache_key = f"{cache_type}_{key}" if key else cache_type
        last_update = self.config["last_update"].get(cache_key)
        
        if not last_update:
            return False
        
        # 计算缓存是否过期
        last_update_time = datetime.fromisoformat(last_update)
        ttl = self.config.get(f"{cache_type}_ttl", 3600)  # 默认缓存1小时
        
        return datetime.now() - last_update_time < timedelta(seconds=ttl)
    
    def get_cache(self, cache_type, key=None):
        """获取缓存数据"""
        if not self.is_cache_valid(cache_type, key):
            return None
        
        cache_path = self.get_cache_path(cache_type, key)
        try:
            with open(cache_path, 'rb') as f:
                return pickle.load(f)
        except Exception as e:
            print(f"读取缓存出错: {str(e)}")
            return None
    
    def set_cache(self, data, cache_type, key=None):
        """设置缓存数据"""
        cache_path = self.get_cache_path(cache_type, key)
        try:
            with open(cache_path, 'wb') as f:
                pickle.dump(data, f)
            
            # 更新最后更新时间
            cache_key = f"{cache_type}_{key}" if key else cache_type
            self.config["last_update"][cache_key] = datetime.now().isoformat()
            self.save_config()
            
            return True
        except Exception as e:
            print(f"保存缓存出错: {str(e)}")
            return False
    
    def clear_cache(self, cache_type=None, key=None):
        """清除缓存数据"""
        if cache_type:
            if key:
                # 清除特定缓存
                cache_path = self.get_cache_path(cache_type, key)
                if os.path.exists(cache_path):
                    os.remove(cache_path)
                
                # 更新配置
                cache_key = f"{cache_type}_{key}"
                if cache_key in self.config["last_update"]:
                    del self.config["last_update"][cache_key]
            else:
                # 清除特定类型的所有缓存
                for file in os.listdir(self.cache_dir):
                    if file.startswith(f"{cache_type}_") and file.endswith(".pkl"):
                        os.remove(os.path.join(self.cache_dir, file))
                
                # 更新配置
                keys_to_remove = []
                for key in self.config["last_update"]:
                    if key.startswith(f"{cache_type}_"):
                        keys_to_remove.append(key)
                
                for key in keys_to_remove:
                    del self.config["last_update"][key]
        else:
            # 清除所有缓存
            for file in os.listdir(self.cache_dir):
                if file.endswith(".pkl"):
                    os.remove(os.path.join(self.cache_dir, file))
            
            # 重置配置
            self.config["last_update"] = {}
        
        self.save_config()
